#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
LangGraph 简单入门示例
这个示例不需要API密钥，可以直接运行
"""

from typing import List, Dict, Any
from langchain_core.tools import tool
from langchain_core.messages import HumanMessage, AIMessage
import json

# 定义工具函数
@tool
def get_weather(city: str) -> str:
    """获取指定城市的天气信息"""
    # 模拟天气数据
    weather_data = {
        "北京": "晴天，温度25°C，湿度60%",
        "上海": "多云，温度22°C，湿度70%", 
        "广州": "雨天，温度28°C，湿度85%",
        "深圳": "晴天，温度30°C，湿度65%",
        "杭州": "阴天，温度20°C，湿度75%",
        "成都": "多云，温度24°C，湿度68%",
        "beijing": "晴天，温度25°C，湿度60%",
        "shanghai": "多云，温度22°C，湿度70%"
    }
    return weather_data.get(city, f"{city}的天气：晴天，温度适宜，湿度正常")

@tool
def get_time() -> str:
    """获取当前时间"""
    import datetime
    now = datetime.datetime.now()
    return f"当前时间：{now.strftime('%Y年%m月%d日 %H:%M:%S')}"

@tool
def calculate(expression: str) -> str:
    """计算简单的数学表达式"""
    try:
        # 安全的数学计算
        allowed_chars = set('0123456789+-*/.() ')
        if not all(c in allowed_chars for c in expression):
            return "错误：只支持基本的数学运算"
        
        result = eval(expression)
        return f"计算结果：{expression} = {result}"
    except Exception as e:
        return f"计算错误：{str(e)}"

# 简单的模拟Agent类
class SimpleWeatherAgent:
    def __init__(self):
        self.tools = {
            "get_weather": get_weather,
            "get_time": get_time,
            "calculate": calculate
        }
        self.conversation_history = []
    
    def _analyze_user_input(self, user_input: str) -> Dict[str, Any]:
        """分析用户输入，决定使用哪个工具"""
        user_input_lower = user_input.lower()
        
        # 检查是否询问天气
        if any(keyword in user_input_lower for keyword in ["天气", "weather", "温度", "下雨", "晴天"]):
            # 尝试提取城市名
            cities = ["北京", "上海", "广州", "深圳", "杭州", "成都"]
            for city in cities:
                if city in user_input:
                    return {"tool": "get_weather", "params": {"city": city}}
            # 如果没有找到具体城市，返回默认
            return {"tool": "get_weather", "params": {"city": "北京"}}
        
        # 检查是否询问时间
        elif any(keyword in user_input_lower for keyword in ["时间", "几点", "现在", "time"]):
            return {"tool": "get_time", "params": {}}
        
        # 检查是否要计算
        elif any(keyword in user_input_lower for keyword in ["计算", "算", "+"]) or any(char in user_input for char in "+-*/"):
            # 提取数学表达式
            import re
            math_pattern = r'[\d+\-*/\.\(\)\s]+'
            match = re.search(math_pattern, user_input)
            if match:
                expression = match.group().strip()
                return {"tool": "calculate", "params": {"expression": expression}}
        
        # 默认情况
        return {"tool": "chat", "params": {"message": user_input}}
    
    def _execute_tool(self, tool_name: str, params: Dict[str, Any]) -> str:
        """执行工具"""
        if tool_name in self.tools:
            tool_func = self.tools[tool_name]
            try:
                return tool_func.invoke(params)
            except Exception as e:
                return f"工具执行错误：{str(e)}"
        else:
            return "抱歉，我还不知道如何处理这个请求。我可以帮您查询天气、时间或进行简单计算。"
    
    def _generate_response(self, user_input: str, tool_result: str) -> str:
        """生成最终回复"""
        user_input_lower = user_input.lower()
        
        if "天气" in user_input_lower or "weather" in user_input_lower:
            return f"根据您的询问，我查询到了以下信息：\n{tool_result}"
        elif "时间" in user_input_lower or "time" in user_input_lower:
            return f"为您查询时间：\n{tool_result}"
        elif "计算" in user_input_lower or any(char in user_input for char in "+-*/"):
            return f"计算完成：\n{tool_result}"
        else:
            return f"我是一个简单的助手，{tool_result}"
    
    def chat(self, user_input: str) -> str:
        """主要的对话函数"""
        print(f"用户输入：{user_input}")
        print("正在分析...")
        
        # 分析用户输入
        analysis = self._analyze_user_input(user_input)
        print(f"分析结果：使用工具 {analysis['tool']}")
        
        # 执行工具
        if analysis['tool'] == 'chat':
            response = analysis['params']['message']
            final_response = "我是一个简单的智能助手。我可以：\n1. 查询天气（如：北京天气怎么样？）\n2. 查询时间（如：现在几点？）\n3. 进行计算（如：计算 2+3*4）\n请告诉我您需要什么帮助！"
        else:
            tool_result = self._execute_tool(analysis['tool'], analysis['params'])
            print(f"工具结果：{tool_result}")
            final_response = self._generate_response(user_input, tool_result)
        
        # 保存对话历史
        self.conversation_history.append({
            "user": user_input,
            "assistant": final_response,
            "tool_used": analysis['tool']
        })
        
        return final_response

def main():
    """主函数"""
    print("=" * 50)
    print("🌤️  LangGraph 简单天气助手演示")
    print("=" * 50)
    print("这是一个简单的Agent示例，展示了LangGraph的基本概念")
    print("支持的功能：")
    print("1. 查询天气：'北京天气怎么样？'")
    print("2. 查询时间：'现在几点？'")
    print("3. 简单计算：'计算 2+3*4'")
    print("4. 输入 'quit' 或 'exit' 退出")
    print("=" * 50)
    
    agent = SimpleWeatherAgent()
    
    # 演示一些示例
    print("\n📝 演示示例：")
    demo_inputs = [
        "北京今天天气怎么样？",
        "现在几点了？",
        "计算 15 + 25 * 2"
    ]
    
    for demo_input in demo_inputs:
        print(f"\n示例 -> {demo_input}")
        response = agent.chat(demo_input)
        print(f"回复 -> {response}")
        print("-" * 30)
    
    print("\n🎯 现在轮到您了！请输入您的问题：")
    
    while True:
        try:
            user_input = input("\n您：").strip()
            
            if user_input.lower() in ['quit', 'exit', '退出', 'bye']:
                print("再见！感谢使用LangGraph助手！")
                break
            
            if not user_input:
                print("请输入有效的问题。")
                continue
            
            response = agent.chat(user_input)
            print(f"助手：{response}")
            
        except KeyboardInterrupt:
            print("\n\n程序已终止。再见！")
            break
        except Exception as e:
            print(f"发生错误：{str(e)}")
            print("请重试...")

if __name__ == "__main__":
    main() 